# backend/modules/visualization/routes.py

import os  # 导入操作系统相关功能
from flask_restful import Resource, reqparse  # 导入Flask-RESTful库中的Resource和reqparse类
from flask import url_for, current_app  # 导入Flask库中的url_for和current_app
from .models import Visualization  # 导入当前模块中的Visualization模型
from ...modules.forecasting.models import Forecast  # 确保导入Forecast模型
from flask_jwt_extended import jwt_required, get_jwt_identity  # 导入JWT身份验证相关的函数
from datetime import datetime, timedelta  # 导入日期和时间处理相关功能
import uuid  # 导入UUID生成器
import matplotlib.pyplot as plt  # 导入matplotlib绘图库
import json  # 导入JSON处理库
import pytz  # 导入pytz库，用于处理时区
import matplotlib.dates as mdates  # 导入日期处理模块用于绘图
import pandas as pd  # 导入pandas数据处理库
from ..dataimport.models import UploadedFile  # 导入上传文件模型
from ...extensions import db  # 导入数据库扩展
from ...modules.authentication.models import User  # 导入用户模型


class VisualizationRoutes:
    def __init__(self, api):
        # 初始化可视化路由，将不同的API资源添加到指定路径
        api.add_resource(CreateChart, '/api/visualization/create')  # 创建图表的端点
        api.add_resource(GetCharts, '/api/visualization')  # 获取图表列表的端点
        api.add_resource(DeleteChart, '/api/visualization/delete/<int:chart_id>')  # 删除图表的端点


class CreateChart(Resource):
    @jwt_required()  # JWT身份验证装饰器，确保请求已验证
    def post(self):
        """
        创建图表的API端点。
        需要提供 'chart_type' 和 'forecast_id'。
        """
        current_user = get_current_user()  # 获取当前用户
        if not current_user:
            return {"message": "用户身份验证失败"}, 401  # 如果身份验证失败，返回401状态码

        # 权限控制，检查用户角色是否为管理员或普通用户
        if current_user.role not in ['管理员', '普通用户']:
            return {"message": "未授权访问"}, 403  # 非法访问返回403状态码

        # 创建请求解析器
        parser = reqparse.RequestParser()
        parser.add_argument('chart_type', type=str, required=True, location='json', help="图表类型是必需的")  # 解析图表类型
        parser.add_argument('forecast_id', type=int, required=True, location='json', help="预测ID是必需的")  # 解析预测ID
        args = parser.parse_args()  # 解析参数

        chart_type = args['chart_type']  # 获取图表类型
        forecast_id = args['forecast_id']  # 获取预测ID

        # 验证图表类型是否合法
        ALLOWED_CHART_TYPES = {'line', 'bar', 'pie'}  # 允许的图表类型
        if chart_type not in ALLOWED_CHART_TYPES:
            return {"message": "不支持的图表类型"}, 400  # 不支持的图表类型返回400状态码

        # 获取预测记录
        forecast = Forecast.query.get(forecast_id)  # 查询数据库中的预测记录
        if not forecast:
            return {"message": "预测记录不存在"}, 404  # 预测记录不存在返回404状态码

        # 权限控制，确保用户只能访问自己的预测记录
        if current_user.role != '管理员' and forecast.uploaded_file.user_id != current_user.id:
            return {"message": "未授权生成其他用户的图表"}, 403  # 非管理员且不是该预测记录的拥有者返回403状态码

        try:
            # 获取预测结果
            prediction_list = json.loads(forecast.result)  # 将预测结果从JSON格式加载
            if not isinstance(prediction_list, list):
                return {"message": "预测结果格式不正确"}, 400  # 如果预测结果格式不正确返回400状态码

            # 获取最后日期并转换为datetime对象（北京时间）
            last_date = forecast.last_date
            if not last_date:
                return {"message": "预测记录缺少最后一个日期信息"}, 400  # 缺少最后一个日期信息返回400状态码

            # 确保last_date是datetime.date类型
            if isinstance(last_date, str):
                last_date = datetime.strptime(last_date, '%Y-%m-%d').date()  # 字符串转换为日期
            elif isinstance(last_date, datetime):
                last_date = last_date.date()  # 已经是datetime类型，直接转换
            elif isinstance(last_date, pd.Timestamp):
                last_date = last_date.to_pydatetime().date()  # pandas时间戳转换为日期

            # 设置北京时间
            beijing_tz = pytz.timezone('Asia/Shanghai')  # 设置北京时区
            # 将last_date转换为datetime对象并设置为北京时间
            start_date = datetime.combine(last_date, datetime.min.time()).replace(tzinfo=beijing_tz)

            # 生成未来的日期列表
            x_dates = [start_date + timedelta(days=i + 1) for i in range(len(prediction_list))]  # 生成日期列表

            # 添加调试日志
            current_app.logger.debug(f"Forecast ID: {forecast.id}, Last Date: {last_date}")
            current_app.logger.debug(f"start_date: {start_date}")
            current_app.logger.debug(f"x_dates sample: {x_dates[:5]} ... {x_dates[-5:]}")
            current_app.logger.debug(f"prediction_list: {prediction_list[:5]} ... {prediction_list[-5:]}")

            # 创建图表
            plt.figure(figsize=(12, 6))  # 设置图表尺寸为12x6英寸

            # 根据图表类型生成图表
            if chart_type == 'line':
                plt.plot(x_dates, prediction_list, marker='o', linestyle='-', color='skyblue', label='需求预测')  # 画线图
                plt.xlabel('日期', fontproperties='SimHei')  # 设置x轴标签
                plt.ylabel('预测数量', fontproperties='SimHei')  # 设置y轴标签
                plt.title(f'{forecast.category} 需求预测趋势图', fontproperties='SimHei')  # 设置标题
                plt.legend(prop={'family': 'SimHei'})  # 设置图例
            elif chart_type == 'bar':
                plt.bar(x_dates, prediction_list, color='lightgreen', label='需求预测')  # 画柱状图
                plt.xlabel('日期', fontproperties='SimHei')  # 设置x轴标签
                plt.ylabel('预测数量', fontproperties='SimHei')  # 设置y轴标签
                plt.title(f'{forecast.category} 需求预测柱状图', fontproperties='SimHei')  # 设置标题
                plt.legend(prop={'family': 'SimHei'})  # 设置图例
            elif chart_type == 'pie':
                # 饼图通常不使用时间轴，这里按月汇总
                df = pd.DataFrame({
                    '日期': x_dates,
                    '预测数量': prediction_list
                })
                df['月份'] = df['日期'].dt.to_period('M')  # 按月份分组
                monthly_data = df.groupby('月份')['预测数量'].sum()  # 按月份聚合数据
                labels = monthly_data.index.strftime('%Y-%m')  # 获取标签
                sizes = monthly_data.values  # 获取对应的值
                plt.pie(sizes, labels=labels, autopct='%1.1f%%', startangle=140)  # 画饼图
                plt.title(f'{forecast.category} 需求预测饼图（按月汇总）', fontproperties='SimHei')  # 设置标题

            # 设置日期格式和布局
            if chart_type in {'line', 'bar'}:
                ax = plt.gca()  # 获取当前坐标轴
                ax.xaxis.set_major_locator(mdates.MonthLocator())  # 设置x轴主刻度为月份
                ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m'))  # 设置x轴日期格式
                plt.xticks(rotation=45, fontproperties='SimHei')  # 设置x轴标签倾斜

            plt.tight_layout()  # 调整布局以适应图表

            # 保存图表为文件
            charts_dir = os.path.join(current_app.root_path, 'static', 'charts')  # 图表保存目录
            os.makedirs(charts_dir, exist_ok=True)  # 创建目录（如果不存在）
            filename = f"{uuid.uuid4().hex}_{chart_type}_{forecast.category}.png"  # 生成唯一文件名
            filepath = os.path.join(charts_dir, filename)  # 生成文件路径
            plt.savefig(filepath, format='png')  # 保存图表
            plt.close()  # 关闭图表

            # 将图表路径保存到数据库
            relative_path = os.path.join('charts', filename).replace('\\', '/')  # 处理文件路径
            chart_entry = Visualization(
                chart_type=chart_type,  # 图表类型
                image_path=relative_path,  # 图表文件路径
                forecast_id=forecast_id,  # 关联的预测ID
                uploaded_file_id=forecast.uploaded_file_id,  # 确保关联上传文件
                created_at=datetime.utcnow()  # 创建时间
            )
            db.session.add(chart_entry)  # 将图表记录添加到数据库
            db.session.commit()  # 提交更改

            # 返回图表的URL
            chart_url = url_for('static', filename=relative_path, _external=True)  # 获取图表的外部URL
            return {"message": "图表创建成功", "chart_id": chart_entry.id, "image_url": chart_url}, 200  # 返回成功消息和图表ID和URL

        except Exception as e:
            db.session.rollback()  # 回滚数据库事务
            current_app.logger.error(f"图表创建失败: {e}", exc_info=True)  # 记录错误日志
            return {"message": f"图表创建失败: {str(e)}"}, 500  # 返回失败消息和500状态码


class GetCharts(Resource):
    @jwt_required()  # JWT身份验证装饰器
    def get(self):
        current_user = get_current_user()  # 获取当前用户
        if not current_user or current_user.role not in ['管理员', '普通用户']:
            return {"message": "未授权访问"}, 403  # 权限控制，未授权返回403状态码

        try:
            if current_user.role == '管理员':
                # 管理员可以查看所有图表
                charts = Visualization.query.order_by(Visualization.created_at.desc()).all()  # 获取所有图表
            else:
                # 普通用户只能查看自己关联的图表
                user_file_ids = [file.id for file in current_user.uploaded_files]  # 获取当前用户上传文件的ID
                charts = Visualization.query.filter(Visualization.uploaded_file_id.in_(user_file_ids)).order_by(Visualization.created_at.desc()).all()  # 获取当前用户关联的图表

            chart_list = []  # 存储图表信息的列表
            for chart in charts:
                forecast = Forecast.query.get(chart.forecast_id)  # 获取关联的预测记录
                if not forecast:
                    forecast_info = {
                        "id": "未知",  # 关联预测的ID
                        "timestamp": "未知",  # 时间戳
                        "model_type": "未知",  # 模型类型
                        "category": "未知"  # 类别
                    }
                else:
                    forecast_info = {
                        "id": forecast.id,  # 关联预测的ID
                        "timestamp": forecast.timestamp.strftime('%Y-%m-%d %H:%M:%S') if forecast.timestamp else "未知",  # 格式化时间戳
                        "model_type": forecast.model_type,  # 模型类型
                        "category": forecast.category if forecast.category else "未知"  # 类别
                    }

                chart_list.append({
                    "id": chart.id,  # 图表ID
                    "type": chart.chart_type,  # 图表类型
                    "image_url": url_for('static', filename=chart.image_path, _external=True),  # 图表URL
                    "created_at": chart.created_at.strftime('%Y-%m-%d %H:%M:%S'),  # 创建时间
                    "forecast_id": chart.forecast_id,  # 预测ID
                    "forecast_info": forecast_info  # 预测信息
                })
            return {"charts": chart_list}, 200  # 返回图表列表和200状态码

        except Exception as e:
            current_app.logger.error(f"获取图表失败: {e}")  # 记录错误
            return {"message": "获取图表失败，请重试。"}, 500  # 返回失败消息和500状态码


class DeleteChart(Resource):
    @jwt_required()  # JWT身份验证装饰器
    def delete(self, chart_id):
        current_user = get_current_user()  # 获取当前用户
        if not current_user:
            return {"message": "用户身份验证失败"}, 401  # 身份验证失败返回401状态码

        try:
            chart = Visualization.query.get(chart_id)  # 获取要删除的图表
            if not chart:
                return {"message": "图表不存在"}, 404  # 图表不存在返回404状态码

            # 获取关联的预测记录
            forecast = Forecast.query.get(chart.forecast_id)
            if not forecast:
                return {"message": "关联的预测记录不存在"}, 404  # 关联的预测记录不存在返回404状态码

            # 获取关联的上传文件
            uploaded_file = UploadedFile.query.get(forecast.uploaded_file_id)
            if not uploaded_file:
                return {"message": "关联的上传文件不存在"}, 404  # 关联的上传文件不存在返回404状态码

            # 获取图表的上传用户
            target_user = uploaded_file.user if uploaded_file.user else None
            if not target_user:
                return {"message": "关联的用户不存在"}, 404  # 关联的用户不存在返回404状态码

            # 权限控制
            if current_user.role == '管理员':
                if target_user.role == '管理员' and not current_user.is_pzy():
                    return {"message": "您无权删除其他管理员的图表"}, 403  # 管理员不能删除其他管理员的图表
                # 管理员可以删除普通用户或pzy管理员的数据
            elif target_user.id == current_user.id:
                # 普通用户可以删除自己的图表
                pass
            else:
                return {"message": "未授权删除该图表"}, 403  # 非法删除返回403状态码

            # 删除图表文件
            chart_file_path = os.path.join(current_app.root_path, 'static', chart.image_path)  # 图表文件路径
            if os.path.exists(chart_file_path):
                os.remove(chart_file_path)  # 删除图表文件
                current_app.logger.info(f"已删除图表文件: {chart_file_path}")  # 记录删除日志
            else:
                current_app.logger.warning(f"图表文件不存在: {chart_file_path}")  # 记录警告日志

            # 从数据库中删除记录
            db.session.delete(chart)  # 删除记录
            db.session.commit()  # 提交更改

            return {"message": "图表删除成功"}, 200  # 返回成功消息和200状态码

        except Exception as e:
            db.session.rollback()  # 回滚数据库事务
            current_app.logger.error(f"图表删除失败: {e}", exc_info=True)  # 记录错误
            return {"message": "图表删除失败，请重试。"}, 500  # 返回失败消息和500状态码


def get_current_user():
    """
    Helper function to retrieve the current user object from JWT.
    """
    current_identity = get_jwt_identity()  # 获取JWT身份
    if not current_identity:
        return None  # 身份不存在返回None
    # 提取用户名
    if isinstance(current_identity, dict):
        username = current_identity.get('username')  # 从字典中获取用户名
        if not username:
            current_app.logger.error("JWT identity does not contain 'username'")
            return None  # 如果用户名不存在返回None
    else:
        username = current_identity  # 直接使用身份作为用户名

    user = User.query.filter_by(username=username).first()  # 根据用户名查询用户
    return user  # 返回用户对象
